#include "wrapper_sstructmg.h"

const idx_t glb_nparts = 4;
const idx_t num_diag = 19;

int main(int argc, char* argv[])
{
    int my_pid, num_procs;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    MPI_Comm_size(MPI_COMM_WORLD, &num_procs);
    if (my_pid == 0) printf("\033[1;35mMax threads %d\033[0m \n", omp_get_max_threads());

    idx_t arg_cnt = 1;
    const idx_t glb_dims [3] = { atoi(argv[arg_cnt++]), atoi(argv[arg_cnt++]), atoi(argv[arg_cnt++]) };
    const idx_t prc_dims [3] = { atoi(argv[arg_cnt++]), atoi(argv[arg_cnt++]), atoi(argv[arg_cnt++]) };// 注意这是块内的进程数
    assert(glb_dims[1] % glb_nparts == 0);
    const idx_t blk_dims[3] = { glb_dims[0], glb_dims[1] / glb_nparts, glb_dims[2] };

    assert(prc_dims[0] * prc_dims[1] * prc_dims[2] * glb_nparts == num_procs);
    assert(prc_dims[2] == 1);
    assert(blk_dims[0] % prc_dims[0] == 0 && blk_dims[1] % prc_dims[1] == 0);

    const std::string pathname(argv[arg_cnt++]);
    const std::string   its_name = std::string(argv[arg_cnt++]),
                        prc_name = std::string(argv[arg_cnt++]);
    std::string config_mg_file;
	if (prc_name == "GMG") config_mg_file = std::string(argv[arg_cnt++]);


    const int nprocs_per_part = prc_dims[0] * prc_dims[1] * prc_dims[2];
    const idx_t my_part = my_pid / nprocs_per_part,
                pid_in_part = my_pid - my_part * nprocs_per_part;
    idx_t cart_ids[3];
    {// 计算本进程在part内处于什么位置
        cart_ids[0] =  pid_in_part / (prc_dims[1] * prc_dims[2]);
        cart_ids[1] = (pid_in_part - cart_ids[0] * prc_dims[1] * prc_dims[2]) / prc_dims[2];
        cart_ids[2] = (pid_in_part - cart_ids[0] * prc_dims[1] * prc_dims[2] - cart_ids[1] * prc_dims[2]);
    }

    idx_t my_ilower[3], my_iupper[3], box_ends[3], my_dims[3];
    idx_t my_nelems = 1;
    for (idx_t d = 0; d < 3; d++) {
        my_ilower[d] =  cart_ids[d]     * (blk_dims[d] / prc_dims[d]);
        my_iupper[d] = (cart_ids[d] + 1)* (blk_dims[d] / prc_dims[d]) - 1;
        box_ends[d] = my_iupper[d] + 1;
        my_dims[d] = my_iupper[d] - my_ilower[d] + 1;
        my_nelems *= my_dims[d];
    }
#ifdef DEBUG_PRINT
    for (int p = 0; p < num_procs; p++) {
        if (my_pid == p) {
            printf("Proc %d part %d: (%d,%d,%d)~(%d,%d,%d)\n", my_pid, my_part,
                my_ilower[0], my_ilower[1], my_ilower[2],
                my_iupper[0], my_iupper[1], my_iupper[2] );
            fflush(stdout);
        }
        MPI_Barrier(MPI_COMM_WORLD);
    }
    MPI_Barrier(MPI_COMM_WORLD);
#endif

    // 先读入本进程所负责的区域的数据
    data_t * A_buf = new data_t [my_nelems * num_diag];
    data_t * rhs_buf = new data_t [my_nelems];
    data_t * x0_buf  = new data_t [my_nelems];

    MPI_Barrier(MPI_COMM_WORLD);
    if (my_pid == 0) printf("begin reading data\n");

    {// 读入
        float * vec_buf = new float [my_nelems];

        MPI_File fh = MPI_FILE_NULL;// 文件句柄
        const MPI_Datatype etype = MPI_FLOAT;
        MPI_Datatype read_type = MPI_DATATYPE_NULL;// 写出类型
        idx_t size[3], subsize[3], start[3];
        size   [0] = glb_dims [0]; size   [1] = glb_dims [1]; size   [2] = glb_dims [2];
        subsize[0] = my_dims  [0]; subsize[1] =  my_dims [1]; subsize[2] =  my_dims [2];
        start  [0] = my_ilower[0]; start  [1] = my_ilower[1] + my_part * glb_dims[1] / glb_nparts;// 注意这里
        start  [2] = my_ilower[2];

        MPI_Type_create_subarray(3, size, subsize, start, MPI_ORDER_C, etype, &read_type);
        MPI_Type_commit(&read_type);
        MPI_Offset displacement = 0;
        displacement *= sizeof(*vec_buf);// 位移要以字节为单位！
        MPI_Status status;

        int ret;
        ret = MPI_File_open(MPI_COMM_WORLD, (pathname + "/array_b").c_str(), MPI_MODE_RDONLY, MPI_INFO_NULL, &fh);
        if (ret != MPI_SUCCESS) {
            printf("Could not open b file: ret %d\n", ret);
        }
        MPI_File_set_view(fh, displacement, etype, read_type, "native", MPI_INFO_NULL);
        MPI_File_read_all(fh, vec_buf, my_nelems, etype, &status);
        MPI_File_close(&fh);
        #pragma omp parallel for schedule(static)
        for (idx_t e = 0; e < my_nelems; e++)
            rhs_buf[e] = vec_buf[e];

        ret = MPI_File_open(MPI_COMM_WORLD, (pathname + "/array_x").c_str(), MPI_MODE_RDONLY, MPI_INFO_NULL, &fh);
        if (ret != MPI_SUCCESS) {
            printf("Could not open x file: ret %d\n", ret);
        }
        MPI_File_set_view(fh, displacement, etype, read_type, "native", MPI_INFO_NULL);
        MPI_File_read_all(fh, vec_buf, my_nelems, etype, &status);
        MPI_File_close(&fh);
        #pragma omp parallel for schedule(static)
        for (idx_t e = 0; e < my_nelems; e++)
            x0_buf[e] = vec_buf[e];
        
        // 读入SOA存储的矩阵
        for (idx_t d = 0; d < num_diag; d++) {// 逐条对角线
            ret = MPI_File_open(MPI_COMM_WORLD, (pathname + "/array_a." + std::to_string(d)).c_str(), MPI_MODE_RDONLY, MPI_INFO_NULL, &fh);
            if (ret != MPI_SUCCESS) {
                printf("Could not open x file: ret %d\n", ret);
            }
            MPI_File_set_view(fh, displacement, etype, read_type, "native", MPI_INFO_NULL);
            MPI_File_read_all(fh, vec_buf, my_nelems, etype, &status);
            MPI_File_close(&fh);

            #pragma omp parallel for schedule(static)
            for (idx_t e = 0; e < my_nelems; e++)
                A_buf[e * num_diag + d] = vec_buf[e];
        }
        delete [] vec_buf;

        MPI_Type_free(&read_type);
    }
    
    MPI_Barrier(MPI_COMM_WORLD);
    if (my_pid == 0) printf("done reading data\n");

    std::vector< Entry<idx_t, data_t> > all_nnz;
    {// 提取块边界非零元
            #pragma omp parallel for collapse(2) schedule(static)
            for (idx_t i0 = my_ilower[0]; i0 <= my_iupper[0]; i0++)
            for (idx_t i1 = my_ilower[1]; i1 <= my_iupper[1]; i1++)
            for (idx_t i2 = my_ilower[2]; i2 <= my_iupper[2]; i2++) {
                StructDescriptor<idx_t> row, col;
                data_t* dst_ptr = A_buf + (((i0 - my_ilower[0] ) * (my_iupper[1] - my_ilower[1] + 1)
                                        +    i1 - my_ilower[1] ) * (my_iupper[2] - my_ilower[2] + 1)
                                        +    i2 - my_ilower[2] ) * num_diag;
                row.id = my_part;
                row.index[0] = i0; row.index[1] = i1; row.index[2] = i2;
                
                for (idx_t d = 0; d < num_diag; d++) {
                    const idx_t * ngb_offsets = stencil_offset_3d19 + d * 3;
                    idx_t ngb_i[3] = {  i0 + ngb_offsets[0],
                                        i1 + ngb_offsets[1],
                                        i2 + ngb_offsets[2] };
                    if (0 <= ngb_i[0] && ngb_i[0] < blk_dims[0] &&
                        0 <= ngb_i[1] && ngb_i[1] < blk_dims[1] ) {
                        if (ngb_i[2] < 0 || ngb_i[2] >= blk_dims[2]) {
                            assert(dst_ptr[d] == 0.0);
                        }
                        continue;
                    } else {
                        if (ngb_i[2] < 0 || ngb_i[2] >= blk_dims[2]) {
                            assert(dst_ptr[d] == 0.0);
                            continue;
                        }

                        if (0 <= ngb_i[0] && ngb_i[0] < blk_dims[0]) {// 无需跨极点，普通的周期边界
                            // 则必然是ngb_i[1]超出范围了
                            if (ngb_i[1] < 0) { assert(ngb_i[1] == -1);
                                col.id = (my_part - 1 >= 0) ? (my_part - 1) : glb_nparts - 1;
                                ngb_i[1] = blk_dims[1] - 1;
                            } else { assert(ngb_i[1] == blk_dims[1]);
                                col.id = (my_part + 1 < glb_nparts) ? (my_part + 1) : 0;
                                ngb_i[1] = 0;
                            }
                        } else {// 跨极点
                            idx_t oppo_part = (my_part + glb_nparts / 2) % glb_nparts;// 正对面的块
                            if (ngb_i[0] < 0) { assert(ngb_i[0] == -1);
                                ngb_i[0] = 0;
                            } else { assert(ngb_i[0] == blk_dims[0]);
                                ngb_i[0] = blk_dims[0] - 1;
                            }
                            if (ngb_i[1] < 0) { assert(ngb_i[1] == -1);
                                col.id = (oppo_part - 1 >= 0) ? (oppo_part - 1) : glb_nparts - 1;
                                ngb_i[1] = blk_dims[1] - 1;
                            } else if (ngb_i[1] >= blk_dims[1]) { assert(ngb_i[1] == blk_dims[1]);
                                col.id = (oppo_part + 1 < glb_nparts) ? (oppo_part + 1) : 0;
                                ngb_i[1] = 0;
                            } else {
                                col.id = oppo_part;
                            }
                        }

                        col.index[0] = ngb_i[0]; assert(ngb_i[0] >= 0 && ngb_i[0] < blk_dims[0]);
                        col.index[1] = ngb_i[1]; assert(ngb_i[1] >= 0 && ngb_i[1] < blk_dims[1]);
                        col.index[2] = ngb_i[2]; assert(ngb_i[2] >= 0 && ngb_i[2] < blk_dims[2]);
                        all_nnz.push_back(Entry<idx_t, data_t>(row, col, dst_ptr[d]));
                        dst_ptr[d] = 0.0;
                    }
                }
            }
    }
    
    MPI_Barrier(MPI_COMM_WORLD);
    if (my_pid == 0) printf("extract boundary coeffs\n");

    SStructMG_Init();
    {
        SStructGrid<idx_t> ssgrid(MPI_COMM_WORLD, 3, glb_nparts);
        const idx_t halos[3] = {1, 1, 1};
        {// 添加本进程所拥有的块和盒子
            const bool periods[3] = {false, false, false};
            ssgrid.addBlock(my_part, blk_dims, periods);
            ssgrid.addBox(my_part, my_ilower, box_ends);
        }
        ssgrid.assemble();
        const idx_t my_nblks = ssgrid.my_blk_list.size(); assert(my_nblks == 1);
        idx_t shalos [my_nblks * 3];
        for (idx_t k = 0; k < my_nblks; k++) {
            shalos[k * 3 + 0] = halos[0];
            shalos[k * 3 + 1] = halos[1];
            shalos[k * 3 + 2] = halos[2];
        }

        // 根据半结构网格建立向量
        par_SstructVector<idx_t, data_t> b(ssgrid, shalos, true);
        par_SstructVector<idx_t, data_t> x(b), y(b), x0(b);
        {// 填入向量数据
            b.set_box_values(my_part, my_ilower, box_ends, rhs_buf);
            x0.set_box_values(my_part, my_ilower, box_ends, x0_buf);
            y.set_val(0.0, true);
        }
        
        // 根据半结构网格建立矩阵
        __int128_t masks[my_nblks]; for (idx_t k = 0; k < my_nblks; k++) masks[k] = stencil_mask_3d19;
        par_SstructMatrix<idx_t, data_t, data_t> A(ssgrid, shalos, masks);
        {// 填入矩阵数据
            A.init_offd();// 初始化非对角部分
            A.set_box_values(my_part, my_ilower, box_ends, A_buf);
            // 然后处理非对角快，也即非结构部分
            A.add_entries(all_nnz.size(), all_nnz.data());
        }
        A.assemble();

        const int test_cnt = 10;
        std::vector<TEST_RECORD> records;
        for (int test = 0; test < test_cnt; test++) {
            vec_copy(x0, x);

            double ck = vec_dot<idx_t, data_t, double>(b, b);
            if (my_pid == 0) printf("(  b,   b) = %.15e\n", ck);

            ck = vec_dot<idx_t, data_t, double>(x, x);
            if (my_pid == 0) printf("(  x,   x) = %.15e\n", ck);

            A.Mult(b, y, false);
            ck = vec_dot<idx_t, data_t, double>(y, y);
            if (my_pid == 0) printf("(A*b, A*b) = %.15e\n", ck);

            A.Mult(x, y, false);
            ck = vec_dot<idx_t, data_t, double>(y, y);
            if (my_pid == 0) printf("(A*x, A*x) = %.15e\n", ck);

            TEST_CONFIG config;
            config.my_nblks = my_nblks;
            config.config_mg_file = config_mg_file;
            config.ssgrid_ptr = & ssgrid;
            config.atol = 1.0e-6;
            TEST_RECORD rec;

            buildup_solver(its_name, prc_name, config);
            setup_and_solve(A, b, x, rec);

            data_t true_r_norm, b_norm;
            check_residual(A, x, b, y, true_r_norm, b_norm);

            if (my_pid == 0) {
                printf("\033[1;35mtrue ||r|| = %20.16e ||r||/||b||= %20.16e\033[0m\n", 
                    true_r_norm, true_r_norm / b_norm);
                printf("Proc %d Setup, Solve costs %.6f %.6f s\n",
                    my_pid, rec.setup, rec.solve);
            }

            stat_part_times(rec);
            records.push_back(rec);

            destroy_solver();
        }// test loop

        // 输出最佳时间
        if (my_pid == 0) {
            stat_multi_runs(records);
        }
    }// all Semi-** object must be destroyed before MPI_Finalize()
    SStructMG_Finalize();
    MPI_Finalize();
    return 0;
}
